# 人工智能NLP-Agent数字人项目-04-基金数据问答任务工单V1.1--2.14
import abc
import logging
import sys
import pandas as pd
import sqlite3
from typing import Any
from langchain.tools import BaseTool
from utils.instances import TOKENIZER, LLM
from FinSQL_01_generate import generate_sql
from FinSQL_02_query import query_db
from FinSQL_03_answer_from_SQL import generate_answer
import utils.configFinRAG as configFinRAG

# Set up logging
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logger = logging.getLogger()
logger.addHandler(logging.StreamHandler(stream=sys.stdout))


class FinSQLRAG(BaseTool, abc.ABC):
    name: str = "查询金融数据库"
    description: str = "当被问到金融查询相关的问题时，会去金融数据库检索结果"

    def __init__(self):
        super().__init__()

        # Initialize example data once
        self.g_example_question_list = []
        self.g_example_sql_list = []
        self.g_example_fa_list = []
        self.g_example_info_list = []
        self.g_example_token_list = []

        self._load_example_data()

    def _load_example_data(self):
        """Load and tokenize example data once."""
        try:
            sql_examples_file = pd.read_csv(configFinRAG.sql_examples_path, delimiter=",", header=0)

            # Storing columns into lists
            self.g_example_question_list = sql_examples_file['问题'].tolist()
            self.g_example_sql_list = sql_examples_file['SQL'].tolist()
            self.g_example_info_list = sql_examples_file['资料'].tolist()
            self.g_example_fa_list = sql_examples_file['FA'].tolist()

            # Tokenize all questions once
            self.g_example_token_list = [TOKENIZER(q)['input_ids'] for q in self.g_example_question_list]

            logger.info("Example data loaded and tokenized successfully.")

        except Exception as e:
            logger.error(f"Error loading example data: {e}")
            raise RuntimeError("Failed to load and tokenize example data.")

    def _run(self, para) -> str:
        query = para
        try:
            # Generate the SQL query based on the input query
            result_prompt, sql = generate_sql(query, LLM, self.g_example_question_list,
                                              self.g_example_sql_list, self.g_example_token_list)
            logger.info(f"Generated SQL: {sql}")

            # Execute the query on the database
            exc_result = self._execute_sql_query(sql)

            # Generate the final answer
            answer = generate_answer(query, exc_result, LLM, self.g_example_question_list,
                                     self.g_example_info_list, self.g_example_fa_list,
                                     self.g_example_token_list)
            return answer

        except Exception as e:
            logger.error(f"Error in _run method: {e}")
            return 'FinSQLRAG处理异常！'

    def _execute_sql_query(self, sql: str):
        """Execute SQL query and return the results."""
        try:
            # Using 'with' to manage the SQLite connection context
            with sqlite3.connect(r'C:\Users\赵奇\Desktop\Fay-fay-agent-edition0830\data\博金杯比赛数据.db') as conn:
                cs = conn.cursor()
                success_flag, exc_result = query_db(sql, cs)

            if not success_flag:
                logger.warning("SQL query execution failed.")
            return exc_result
        except sqlite3.Error as e:
            logger.error(f"SQL execution error: {e}")
            raise RuntimeError("Failed to execute SQL query.")

    async def _arun(self, *args: Any, **kwargs: Any) -> Any:
        # Not implemented as per the use case
        pass


if __name__ == "__main__":
    tool = FinSQLRAG()
    result = tool.run(
        "请帮我计算，在20210715，中信行业分类划分的一级行业为消费者服务行业中，涨跌幅最大股票的股票代码是？涨跌幅是多少？百分数保留两位小数。股票涨跌幅定义为：（收盘价 - 前一日收盘价 / 前一日收盘价）* 100%。")
    print(result)

